import base_util
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from modelscope import snapshot_download




class BaseSegmentModel:
    """
    抽象基类，定义了文本分块模型的接口。
    """
    def __init__(self, model_name):
        #
        self.model_download_path=f"{base_util.current_project_dir}\\\models\SegmentModels\\"
        self.model_name=model_name
        self.model=None

    def download_model(self,model_id,model_path):
        model_dir = snapshot_download(model_id, cache_dir=model_path)
        return model_dir


    def load_model(self,model_path):
        pass



class BertModel(BaseSegmentModel):
    """
    不扩展了 目前就只跑一个bert就行了
    """

    def __init__(self,model_name='damo/nlp_bert_document-segmentation_chinese-base'):
        #
        super().__init__(model_name)

        self.model_path = snapshot_download(model_name,
                                            cache_dir=self.model_download_path)

        self.model = pipeline(
            task=Tasks.document_segmentation,
            model=self.model_path)


    def segment(self, text):
        """
        将文本 进行分块
        :param text: 输入文本
        :return:  分块的文本
        """
        result = self.model(documents=text)
        return result[OutputKeys.TEXT].split("\n\n")



if __name__ == '__main__':
    bert_model=BertModel()
    txts = bert_model.segment('移动端语音唤醒模型，检测关键词为“小云小云”。模型主体为4层FSMN结构，使用CTC训练准则，参数量750K，适用于移动端设备运行。模型输入为Fbank特征，输出为基于char建模的中文全集token预测，测试工具根据每一帧的预测数据进行后处理得到输入音频的实时检测结果。模型训练采用“basetrain + finetune”的模式，basetrain过程使用大量内部移动端数据，在此基础上，使用1万条设备端录制安静场景“小云小云”数据进行微调，得到最终面向业务的模型。后续用户可在basetrain模型基础上，使用其他关键词数据进行微调，得到新的语音唤醒模型，但暂时未开放模型finetune功能。')
    for txt in txts:
        print(txt)